import os
import subprocess
import requests
import re
# 在文件头部添加
import psutil
import socket
import platform
import time   
import pyfiglet
import tkinter as tk
from tkinter import filedialog
import webbrowser             
from typing import Optional, Tuple, List, Dict
from colorama import init, Fore, Back, Style
from pyfiglet import Figlet
from pyfiglet import fonts
import json  # 添加这一行
import ctypes
import sys

# 初始化颜色
init(autoreset=True)

# 配置你的 DeepSeek API 密钥
API_KEY = "sk-7815728655fa4a258bca884fb37583c2"
API_URL = "https://api.deepseek.com/v1/chat/completions"

# 全局变量存储对话历史
conversation_history: List[Dict[str, str]] = [
    {
        "role": "system",
        "content": "你是一个专业的windows电脑助手，采用计划模式工作。每次只生成一个步骤的可执行代码（Python或Bash），"
                   "代码必须用```python或```bash标记包围，执行完成后需要明确给出以下判断：\n"
                   "1. 如果任务已完成，回复『计划完成』\n"
                   "2. 如果需要继续执行，回复『继续下一步』\n"
                   "3. 如果执行失败需要修正，回复『需要修正』\n"
                   "4. 这些提示词是你在特定情况下给用户的回复，不是让用户对你说的\n"
                   "严格遵循以下规则：\n"
                   "- 必须检查执行结果中的输出内容是否符合预期，每个对话只能生成一个步骤\n"
                   "- 必须验证代码执行后的系统状态变化\n" 
                   "- 必须确认用户原始需求中的每个要求都已满足\n"
                   "- 计划完成前必须生成验证代码检查结果\n"
                   "- 危险操作必须包含回滚方案\n"
                  "- 对于服务器启动类任务，必须持续监控日志直到出现『Done!』或就绪提示\n"
                   "- 必须检测EULA协议是否同意（检查eula.txt文件）\n"
                   "- 服务器控制台必须有『计划完成』明确标记才能结束任务\n"
                   "- 如果执行结果包含错误信息必须优先处理\n"
    }
]

# 存储所有生成的代码片段
generated_code_history: List[Dict[str, str]] = []

def print_header():
    """打印彩色标题"""
    f = Figlet(font='slant')
    print(Fore.CYAN + f.renderText('DEEP Control '))
    print(Fore.YELLOW + "="*60)
    print(Fore.GREEN + "MADE BY ZHOUYI".center(60))
    print(Fore.YELLOW + "="*60)
    print(Fore.RED + "输入“1”即刻体验".center(60))
    print(Fore.YELLOW + "="*60)
    print(Fore.BLUE +"POWERED BY DEEPSEEK V3".center(60))
    print(Fore.YELLOW + "="*60)
    print(Fore.MAGENTA + "DeepSeekAPI强力驱动".center(60))
    print(Fore.YELLOW + "="*60)

def print_menu():
    """打印主菜单"""
    print(Fore.CYAN + "\n" + "="*30 + " 主菜单 " + "="*30)
    print(Fore.GREEN + "1. 发送指令到AI")
    print(Fore.GREEN + "2. 查看系统信息")
    print(Fore.GREEN + "3. 文件管理工具")
    print(Fore.GREEN + "4. 进程管理工具")
    print(Fore.GREEN + "5. 网络工具")
    print(Fore.GREEN + "6. 查看对话历史")
    print(Fore.GREEN + "7. 查看代码历史")
    print(Fore.GREEN + "8. 清除上下文")
    print(Fore.GREEN + "9. 帮助信息")
    print(Fore.GREEN + "0. 退出程序")
    print(Fore.CYAN + "="*68)

def print_system_info():
    """打印系统信息"""
    print(Fore.CYAN + "\n" + "="*30 + " 系统信息 " + "="*30)
    print(Fore.YELLOW + f"操作系统: {platform.system()} {platform.release()}")
    print(Fore.YELLOW + f"系统版本: {platform.version()}")
    print(Fore.YELLOW + f"计算机名称: {platform.node()}")
    print(Fore.YELLOW + f"处理器: {platform.processor()}")
    print(Fore.YELLOW + f"Python版本: {platform.python_version()}")
    print(Fore.CYAN + "="*68)

def get_ai_response(prompt: str) -> Optional[str]:
    """获取DeepSeek AI的响应，包含上下文，支持流式输出"""
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    # 添加用户新消息到历史记录
    conversation_history.append({"role": "user", "content": prompt})
    
    payload = {
        "model": "deepseek-chat",
        "messages": conversation_history,
        "temperature": 0.3,
        "max_tokens": 2000,
        "stream": True  # 启用流式输出
    }
    
    try:
        print(Fore.YELLOW + "\n正在与AI通信...")
        response = requests.post(API_URL, headers=headers, json=payload, stream=True)
        response.raise_for_status()
        
        ai_message = ""
        print(Fore.CYAN + "\nAI回复:")
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                if decoded_line.startswith('data: '):
                    chunk = decoded_line[6:]
                    if chunk == '[DONE]':
                        break
                    try:
                        data = json.loads(chunk)
                        delta = data.get('choices', [{}])[0].get('delta', {})
                        content = delta.get('content', '')
                        ai_message += content
                        print(content, end='', flush=True)
                    except json.JSONDecodeError:
                        continue
        
        print()  # 换行
        
        # 添加AI回复到历史记录
        conversation_history.append({"role": "assistant", "content": ai_message})
        
        # 限制历史记录长度以避免超出token限制
        if len(conversation_history) > 10:
            conversation_history.pop(1)  # 保留系统消息，移除最早的对话
            
        return ai_message
    except Exception as e:
        print(Fore.RED + f"获取AI响应失败: {str(e)}")
        return None

def extract_code(response: str) -> Tuple[str, str, str]:
    """从AI响应中提取解释、代码和代码类型"""
    explanation = re.sub(r'```(?:python|bash).*?```', '', response, flags=re.DOTALL).strip()
    
    # 匹配Python或Bash代码块
    code_match = re.search(r'```(python|bash)(.*?)```', response, re.DOTALL)
    if code_match:
        code_type = code_match.group(1).lower()
        code = code_match.group(2).strip()
        return explanation, code, code_type
    return explanation, None, None

def is_module_installed(module_name: str) -> bool:
    """检查Python模块是否已安装"""
    try:
        __import__(module_name)
        return True
    except ImportError:
        return False

def install_missing_packages(packages: List[str]) -> None:
    """使用pip安装缺失依赖"""
    print(Fore.CYAN + "\n正在安装依赖...")
    for pkg in packages:
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", pkg],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
            print(Fore.GREEN + f"成功安装: {pkg}")
        except subprocess.CalledProcessError:
            print(Fore.RED + f"安装失败: {pkg}")

def execute_code(code: str, code_type: str) -> Tuple[bool, str]:
    """执行提取的代码"""
    if not code or not code_type:
        return False, "无有效代码"
    
    try:
        # 依赖库检查（现在可以正确调用is_module_installed）
        if code_type == 'python':
            required_imports = re.findall(r'^\s*import (\w+)', code, re.M) 
            required_from_imports = re.findall(r'^\s*from (\w+)', code, re.M)
            missing_imports = [
                lib for lib in required_imports + required_from_imports 
                if not is_module_installed(lib)
            ]
            
            if missing_imports:
                print(Fore.YELLOW + "\n缺少依赖库:")
                print(Fore.WHITE + ", ".join(missing_imports))
                confirm = input(Fore.YELLOW + "是否尝试自动安装? (y/n): ").lower()
                if confirm == 'y':
                    install_missing_packages(missing_imports)
                else:
                    return False, "用户取消依赖安装"
        
        # 危险操作检测
        danger_keywords = {
            'python': ['os.system', 'subprocess.call', 'rm ', 'del ', 'shutil.rmtree'],
            'bash': ['rm -rf', 'dd ', 'mkfs', 'fdisk', 'chmod 777']
        }
        
        if any(keyword in code.lower() for keyword in danger_keywords[code_type]):
            print(Fore.RED + "\n⚠️ 警告: 此代码包含潜在危险操作!")
            confirm = input(Fore.YELLOW + "确定要执行吗? (y/n): ").lower()
            if confirm != 'y':
                return False, "用户取消执行"

        # 执行不同类型代码
        if code_type == 'python':
            from io import StringIO
            import sys
            import traceback  # 新增异常捕获
            
            old_stdout = sys.stdout
            sys.stdout = buffer = StringIO()
            
            try:
                exec(code)
                error_info = ""
            except Exception as e:
                error_info = f"\n异常堆栈：\n{traceback.format_exc()}"
            
            sys.stdout = old_stdout
            output = buffer.getvalue()
            full_output = f"{output}{error_info}"
            return (error_info == ""), f"Python执行详情：\n{full_output}"
            
        elif code_type == 'bash':
            result = subprocess.run(
                code,
                shell=True,
                capture_output=True,
                text=True,
                encoding='utf-8'
            )
            output = f"退出码: {result.returncode}\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
            return result.returncode == 0, f"Bash执行详情：\n{output}"
            
    except Exception as e:
        return False, f"执行错误：{str(e)}"

def check_completion(response: str) -> bool:
    """根据AI响应判断任务是否完成"""
    completion_patterns = [
        r'计划完成', 
        r'任务已完成',
        r'no further action needed',
        r'\[COMPLETED\]',
        r'Done!',  # 新增Minecraft服务器完成标记
        r'For help, type "help"',  # 新增服务器就绪提示
        r'\[需要修正\]'  # 添加需要修正的匹配模式
    ]
    # 反向检查避免提前结束
    negative_patterns = [
        r'需要修正',
        r'Exception',
        r'Error',
        r'Warning',
        r'FAILED'
    ]
    
    # 必须同时满足：有完成标记且没有错误信息
    return (any(re.search(pattern, response, re.IGNORECASE) 
              for pattern in completion_patterns)
            and not any(re.search(npattern, response) 
                      for npattern in negative_patterns))

def clear_context() -> None:
    """清除对话上下文，保留系统消息"""
    global conversation_history
    conversation_history = [conversation_history[0]]
    print(Fore.GREEN + "对话上下文已清除")

def display_history() -> None:
    """显示当前对话历史"""
    print(Fore.CYAN + "\n" + "="*30 + " 对话历史 " + "="*30)
    if len(conversation_history) == 1:
        print(Fore.YELLOW + "暂无对话历史")
        print(Fore.CYAN + "="*68)
        return
    
    for i, message in enumerate(conversation_history[1:]):  # 跳过系统消息
        role = "你" if message["role"] == "user" else "AI"
        color = Fore.BLUE if message["role"] == "user" else Fore.MAGENTA
        print(f"{i+1}. {color}{role}: {message['content'][:100]}...")
    print(Fore.CYAN + "="*68)

def show_generated_code() -> None:
    """显示所有生成的代码历史"""
    print(Fore.CYAN + "\n" + "="*30 + " 代码历史 " + "="*30)
    if not generated_code_history:
        print(Fore.YELLOW + "没有保存的代码历史")
        print(Fore.CYAN + "="*68)
        return
    
    for idx, code_item in enumerate(generated_code_history, 1):
        print(f"\n[{Fore.GREEN}{idx}{Fore.RESET}] {Fore.YELLOW}{code_item['timestamp']}")
        print(f"{Fore.CYAN}{code_item['description']}")
    
    while True:
        try:
            print(Fore.YELLOW + "\n输入要查看的代码编号(或输入0返回): ")
            choice = input(Fore.GREEN + "> ")
            if choice.lower() in ['0', 'exit', 'back']:
                return
            
            code_idx = int(choice)
            if 1 <= code_idx <= len(generated_code_history):
                print(Fore.CYAN + "\n" + "="*50)
                print(f"{Fore.GREEN}代码 [{code_idx}]:")
                print(Fore.WHITE + generated_code_history[code_idx-1]['content'])
                print(Fore.CYAN + "="*50)
                print(Fore.YELLOW + "\n1. 再次执行此代码")
                print(Fore.YELLOW + "2. 返回代码列表")
                print(Fore.YELLOW + "0. 返回主菜单")
                
                sub_choice = input(Fore.GREEN + "> ")
                if sub_choice == "1":
                    print(Fore.YELLOW + "\n正在执行该代码...")
                    execute_code(generated_code_history[code_idx-1]['content'])
                    print(Fore.CYAN + "="*50)
                elif sub_choice == "0":
                    return
                # 其他情况会继续显示代码列表
        except ValueError:
            print(Fore.RED + "请输入有效的数字")

def file_manager():
    """文件管理工具"""
    while True:
        print(Fore.CYAN + "\n" + "="*30 + " 文件管理 " + "="*30)
        print(Fore.GREEN + "1. 列出当前目录内容")
        print(Fore.GREEN + "2. 切换目录")
        print(Fore.GREEN + "3. 创建目录")
        print(Fore.GREEN + "4. 删除文件/目录")
        print(Fore.GREEN + "5. 复制文件")
        print(Fore.GREEN + "6. 移动文件")
        print(Fore.GREEN + "0. 返回主菜单")
        print(Fore.CYAN + "="*68)
        
        choice = input(Fore.GREEN + "> ")
        
        if choice == "1":
            print(Fore.YELLOW + "\n当前目录内容:")
            print(Fore.WHITE + "\n".join(os.listdir()))
        elif choice == "2":
            new_dir = input(Fore.YELLOW + "输入目标目录路径: ")
            try:
                os.chdir(new_dir)
                print(Fore.GREEN + f"已切换到: {os.getcwd()}")
            except Exception as e:
                print(Fore.RED + f"错误: {str(e)}")
        elif choice == "3":
            dir_name = input(Fore.YELLOW + "输入要创建的目录名: ")
            try:
                os.mkdir(dir_name)
                print(Fore.GREEN + f"目录 '{dir_name}' 创建成功")
            except Exception as e:
                print(Fore.RED + f"错误: {str(e)}")
        elif choice == "0":
            return
        else:
            print(Fore.RED + "此功能正在开发中，请使用AI指令功能")

def process_manager():
    """进程管理工具"""
    while True:
        print(Fore.CYAN + "\n" + "="*30 + " 进程管理 " + "="*30)
        print(Fore.GREEN + "1. 列出所有进程")
        print(Fore.GREEN + "2. 结束进程")
        print(Fore.GREEN + "0. 返回主菜单")
        print(Fore.CYAN + "="*68)
        
        choice = input(Fore.GREEN + "> ")
        
        if choice == "1":
            print(Fore.YELLOW + "\n正在获取进程列表...")
            try:
                if platform.system() == "Windows":
                    output = subprocess.check_output("tasklist", shell=True).decode()
                else:
                    output = subprocess.check_output("ps aux", shell=True).decode()
                print(Fore.WHITE + output)
            except Exception as e:
                print(Fore.RED + f"获取进程列表失败: {str(e)}")
        elif choice == "0":
            return
        else:
            print(Fore.RED + "此功能正在开发中，请使用AI指令功能")

def network_tools():
    """网络工具"""
    while True:
        print(Fore.CYAN + "\n" + "="*30 + " 网络工具 " + "="*30)
        print(Fore.GREEN + "1. 测试网络连接")
        print(Fore.GREEN + "2. 获取IP信息")
        print(Fore.GREEN + "0. 返回主菜单")
        print(Fore.CYAN + "="*68)
        
        choice = input(Fore.GREEN + "> ")
        
        if choice == "1":
            host = input(Fore.YELLOW + "输入要测试的主机(默认: 8.8.8.8): ") or "8.8.8.8"
            print(Fore.YELLOW + f"\n正在ping {host}...")
            try:
                if platform.system() == "Windows":
                    output = subprocess.check_output(f"ping {host}", shell=True).decode()
                else:
                    output = subprocess.check_output(f"ping -c 4 {host}", shell=True).decode()
                print(Fore.WHITE + output)
            except Exception as e:
                print(Fore.RED + f"ping失败: {str(e)}")
        elif choice == "0":
            return
        else:
            print(Fore.RED + "此功能正在开发中，请使用AI指令功能")

def show_help() -> None:
    """显示帮助信息"""
    print(Fore.CYAN + "\n" + "="*30 + " 帮助信息 " + "="*30)
    print(Fore.GREEN + "1. 发送指令到AI - 输入自然语言指令，AI会生成代码并执行")
    print(Fore.GREEN + "2. 系统信息 - 查看计算机的硬件和系统信息")
    print(Fore.GREEN + "3. 文件管理 - 基本的文件操作工具")
    print(Fore.GREEN + "4. 进程管理 - 查看和结束进程")
    print(Fore.GREEN + "5. 网络工具 - 基本的网络诊断工具")
    print(Fore.GREEN + "6. 对话历史 - 查看与AI的对话记录")
    print(Fore.GREEN + "7. 代码历史 - 查看所有生成的代码")
    print(Fore.GREEN + "8. 清除上下文 - 重置AI的记忆")
    print(Fore.GREEN + "0. 退出 - 退出程序")
    print(Fore.CYAN + "="*68)

def report_execution_result(success: bool, result: str) -> None:
    """将代码执行结果反馈给AI"""
    status = "成功" if success else "失败"
    feedback = f"""代码执行状态：{status}
执行结果详情：
{'-'*40}
{result}
{'-'*40}

请严格检查：
1. 输出内容是否包含预期结果
2. 是否有未处理的错误信息
3. 是否满足用户原始需求的所有要求
4. 是否需要额外验证步骤

根据以上信息生成："""
    
    conversation_history.append({
        "role": "user",
        "content": feedback
    })

def check_admin():
    """检查并以管理员权限重新启动"""
    try:
        is_admin = ctypes.windll.shell32.IsUserAnAdmin()
    except:
        is_admin = False
        
    if not is_admin:
        ctypes.windll.shell32.ShellExecuteW(
            None, "runas", sys.executable, __file__, None, 1
        )
        sys.exit()

def main():
    check_admin()  # 新增权限检查
    print_header()
    
    while True:
        print_menu()
        choice = input(Fore.GREEN + "请选择操作 (1-9, 0退出): ")
        
        if choice == "1":
            user_input = input(Fore.YELLOW + "\n你想让电脑做什么? > ").strip()
            if not user_input:
                continue
                
            print(Fore.YELLOW + "\n正在处理你的请求...")
            
            # 获取初始AI响应
            ai_response = get_ai_response(user_input)
            if not ai_response:
                print(Fore.RED + "无法获取AI响应，请重试")
                continue
                
            while True:  # 计划模式循环
                explanation, code, code_type = extract_code(ai_response)
                
                print(Fore.CYAN + "\nAI计划:")
                print(Fore.WHITE + explanation)
                
                if check_completion(explanation):
                    print(Fore.GREEN + "\n✅ AI判断任务已完成")
                    break
                    
                if code:
                    print(Fore.YELLOW + f"\n即将执行{code_type.upper()}代码:")
                    print(Fore.WHITE + code)
                    
                    # 自动执行逻辑（移除确认步骤）
                    print(Fore.YELLOW + "\n正在执行代码...")
                    success, result = execute_code(code, code_type)
                    
                    # 反馈结果
                    print(Fore.CYAN + "\n执行结果:")
                    print(Fore.GREEN if success else Fore.RED, result)
                    report_execution_result(success, result)  # 修改这里
                    ai_response = get_ai_response("")
                    if not check_completion(ai_response):
                        print(Fore.YELLOW + "\n🔄 AI要求继续执行下一步...")
                    else: 
                        print(Fore.GREEN + "\n✅ AI判断任务已完成")
                        break
                else:
                    print(Fore.YELLOW + "\nAI未提供可执行代码")
                    break
        
        elif choice == "2":
            print_system_info()
        
        elif choice == "3":
            file_manager()
        
        elif choice == "4":
            process_manager()
        
        elif choice == "5":
            network_tools()
        
        elif choice == "6":
            display_history()
        
        elif choice == "7":
            show_generated_code()
        
        elif choice == "8":
            clear_context()
        
        elif choice == "9":
            show_help()
        
        elif choice == "0":
            print(Fore.CYAN + "\n" + "="*30 + " 再见! " + "="*30)
            print(Fore.GREEN + "感谢使用PC助手，再见！".center(60))
            print(Fore.CYAN + "="*68)
            break
        
        else:
            print(Fore.RED + "无效的选项，请重新输入")

if __name__ == "__main__":
    main()
